import matplotlib.image as mpimg
import xmltodict;
from PIL import Image
import imgaug as ia
import os
import json
import random
import copy
import config


ENCODE_METHOD = 'utf-8'

def get_img_path(id):
    '''
    通过id获取图片的路径
    :param id:
    :return:
    '''
    return config.image_dir + id + '.jpg'

def get_annotation_path(id):
    '''
    通过id获取标注的路径
    :param id:
    :return:
    '''
    return config.annotation_dir + id + '.xml'

def read_file_content(file_path):
    '''
    读取文件的内容
    :param file_path:
    :return:
    '''
    f = open(file_path)
    return f.read()


def write_file_content(content,path):
    make_parent_dirs_if_not_exits(path)
    with open(path,"w") as f:
        f.write(content)



def read_image(img_path):
    '''
    加载图片
    :param img_path:
    :return:
    '''
    return mpimg.imread(img_path)

def xml_to_dict(content):
    '''
    xml转化为dict
    :param content:
    :return:
    '''
    return xmltodict.parse(content)

def dict_to_xml(dict):
    return xmltodict.unparse(dict)


def show_img(img_data):
    '''
    显示图像
    :param img_data: 图片数据 narray类型
    :return:
    '''
    Image.fromarray(img_data).show()

def show_img_with_bbs(img_data,bbs):
    '''
    显示图片和bound boxs
    :param img_data:
    :param bbs:
    :return:
    '''
    bbs_img = bbs.draw_on_image(img_data, thickness=2)
    show_img(bbs_img)

def save_img(img_data,path):
    '''
    保存图片
    :param img_data: 图片数据
    :param path: 保存的路径
    :return:
    '''
    make_parent_dirs_if_not_exits(path)
    image = Image.fromarray(img_data)
    image.save(path)

def make_parent_dirs_if_not_exits(path):
    '''
    如果父目录不存在则创建父目录
    :param path:
    :return:
    '''
    dir = os.path.split(path)[0]
    if not os.path.exists(dir):
        os.makedirs(dir)

def get_dict_object(dict):
    '''
    获取object的值,以数组的形式返回
    :param dict:
    :return:
    '''
    obs = []

    adict = dict['annotation']
    if 'object' in adict:
        object = adict['object']

        if(isinstance(object,list)):
            obs = object
        else:
            obs.append(object)

    return obs;


def objects_to_bbs(objects,shape):
    '''
    将objects转化为BoundingBoxes
    :param objects:
    :param shape:
    :return:
    '''
    bbs = []

    for obj in objects:
        bb = __bndbox_to_BoundingBox(obj['bndbox'])
        bbs.append(bb)

    return ia.BoundingBoxesOnImage(bbs,shape=shape)

def __bndbox_to_BoundingBox(bndbox):
    return ia.BoundingBox(x1=int(bndbox['xmin']), y1=int(bndbox['ymin']), x2=int(bndbox['xmax']), y2=int(bndbox['ymax']))





def bbs_to_objects(bbs,demo_obj):
    '''
    将BoundingBoxes转化为objects
    :param bbs:
    :param demo_obj: 示例的object
    :return:
    '''
    res = []

    for box in bbs.bounding_boxes:
        object = copy.deepcopy(demo_obj)
        object = __update_object_dict_by_box(object,box)
        res.append(object)
    return res


def __update_object_dict_by_box(object,box):
    object['bndbox']['xmin'] = str(int(round(min(box.x1,box.x2))))
    object['bndbox']['xmax'] = str(int(round(max(box.x1,box.x2))))
    object['bndbox']['ymin'] = str(int(round(min(box.y1,box.y2))))
    object['bndbox']['ymax'] = str(int(round(max(box.y1,box.y2))))
    return object


def augmenttation_handle(id,aug,new_id,debug=False,drop_empty_percent=1):
    '''
    图片增强处理
    :param id: 图片的id
    :param aug: 增强效果
    :param new_id: 处理后的图片id
    :param debug: 是否显示结果
    :return:
    '''
    img_path = get_img_path(id)
    xml_path = get_annotation_path(id)
    new_img_path = get_img_path(new_id)
    new_xml_path = get_annotation_path(new_id)

    # 处理图像
    img = read_image(img_path)
    image_aug = aug.augment_images([img])[0]
    if debug :
        print(image_aug.shape)


    # 处理xml
    xml_content = read_file_content(xml_path)

    dic = xml_to_dict(xml_content)
    dic['annotation']['filename'] = new_id + '.jpg'

    new_objects = []
    objects = get_dict_object(dic)
    for cls in config.annotation_classes:          #不同类别的box需要单独处理
        class_objects = __get_objects_of_class(objects,cls)
        if len(class_objects) > 0:

            bbs = objects_to_bbs(class_objects,img.shape,)
            bbs_aug = aug.augment_bounding_boxes([bbs])[0]
            bbs_ok = bbs_aug.remove_out_of_image().cut_out_of_image()      #去掉超出边界的box

            if debug :
                if len(bbs_ok.bounding_boxes) > 0 :
                    show_img_with_bbs(image_aug,bbs_ok)

            new_objects.extend(bbs_to_objects(bbs_ok,class_objects[0]))

    dic['annotation']['object'] = new_objects
    out_xml_content = dict_to_xml(dic)

    if len(new_objects) > 0:
        if not random.random() <= drop_empty_percent:        #按照概率丢掉没有标记的图片
            return False

    #保存到磁盘
    save_img(image_aug,new_img_path)
    write_file_content(out_xml_content,new_xml_path)
    print("Save Image:" + new_img_path)
    return True

def __get_objects_of_class(objects,cls):
    res = []
    for obj in objects:
        if obj['name'] == cls:
            res.append(obj)
    return res


def load_ids():
    '''
    加载文件中的所有id
    :return:
    '''
    with open("ids.json",'r') as json_file:
        return json.load(json_file)

